import os
import pickle
import wave

import numpy as np
import opuslib_next
from pydub import AudioSegment

class Opus_Encode:
    def __init__(self):
        self.channels = 1
        self.sample_rate = 16000
        self.sample_width = 2


        self.opus_channels = 1
        self.opus_sample_rate = 16000
        self.opus_sample_width = 2
        self.opus_frame_time = 60
        # 获取每frame的采样数
        self.frame_num = int(self.opus_sample_rate / 1000 * self.opus_frame_time)
        # 计算每frame的采样字节数
        self.opus_frame_size = self.frame_num * self.opus_channels * self.opus_sample_width





    def audio_to_opus(self,audio_file_path):
        file_type = os.path.splitext(audio_file_path)[1]
        if file_type :
            file_type = file_type.lstrip('.')
        audio = AudioSegment.from_file(audio_file_path, format=file_type)
        audio = audio.set_channels(self.channels).set_frame_rate(16000).set_sample_width(self.sample_width)
        # 计算音频总时长
        duration = len(audio)

        # 获取音频原始数据
        raw_data = audio.raw_data

        # 创建一个opus的encoder
        encoder = opuslib_next.Encoder(self.opus_sample_rate, self.opus_channels,opuslib_next.APPLICATION_AUDIO)

        # 获取每frame的采样数
        frame_num = self.frame_num
        # 计算每frame的采样字节数
        frame_bytes_size = self.opus_frame_size
        # 创建一个数组接受转换后的opus数据
        opus_datas = []
        for i in range(0,len(raw_data),frame_bytes_size):
            # 每frame的数据就是一个chunk
            chunk = raw_data[i:i+frame_bytes_size]

            chunk_len = len(chunk)
            # 如果当前chunk的长度不足一个frame的标准长度，通过是最后一个，那么就填充0
            if chunk_len < frame_bytes_size:
                chunk += b'\x00' * (frame_bytes_size - chunk_len)
            np_frame=np.frombuffer(chunk, dtype=np.int16)
            np_frame_bytes=np_frame.tobytes()

            opus_data=encoder.encode(np_frame_bytes,frame_num)
            opus_datas.append(opus_data)
        return opus_datas
    def save_opus_raw(self,opus_datas,output_path):
        with open(output_path, 'wb') as f:
            pickle.dump(opus_datas, f)
    def load_opus_raw(self,output_path):
        with open(output_path, 'rb') as f:
            opus_datas = pickle.load(f)
        return opus_datas
    def opus_to_wav(self,opus_datas):
        output_file = "test.wav"
        decoder = opuslib_next.Decoder(self.opus_sample_rate,self.opus_channels)
        pcm_data = []
        for frame in opus_datas:
            try:
                pcm_frame = decoder.decode(frame,self.opus_frame_size)
                pcm_data.append(pcm_frame)
            except opuslib_next.OpusError as e:
                print(f"解码失败:{e}")
        with wave.open(output_file, 'wb') as w:
            w.setnchannels(self.opus_channels)
            w.setsampwidth(self.opus_sample_width)
            w.setframerate(self.opus_sample_rate)
            print(f"PCM数据：{pcm_data}")
            w.writeframes(b''.join(pcm_data))

if __name__ == '__main__':
    opus = Opus_Encode()
    opus_datas =opus.audio_to_opus('test.mp3')
    opus.save_opus_raw(opus_datas,'test.opus')
    opus_datas_raw=opus.load_opus_raw('test.opus')
    opus.opus_to_wav(opus_datas_raw)


